import hilo
from distorm3 import Decode, Decode32Bits

# 
# Represents an open and attached process.
# Provides some operations:
#   memory derefencing (derefence)
#   making anonymous functions for remote calls (remoteFunction)
# Requires (kind of) that detach() is called before exiting
# Static methods:
#   getRunningProcesses( )
#      Returns a list of running processes
#
class Process:
    def __init__( self, name = None ):
        self.pid = 0 # process id
        self.hProc = None # handle to the (open)
        self.open = False
        if ( name != None ):
            self.attach( name )
    def attach( self, name ):
        if ( self.open ):
            raise Exception()
        self.pid = hilo.getPid( name )
        self.hProc = hilo.openProcess( self.pid )
        self.open = True
    def detach( self ):
        if ( self.open ):
            hilo.closeHandle( self.hProc )
            self.pid = 0
            self.hProc = None
            self.open = False
    def writeInt( self, address, length, val ):
        if ( self.open ):
            hilo.writeInt( self.hProc, address, length, val )
    def readInt( self, address, length ):
        if ( self.open ):
            return hilo.readInt( self.hProc, address, length )
    # todo: readInt
    def isAttached( self ):
        return self.open
    def derefence( self, type, address ):
        pass
    def genCallbackTrace( self, code, grabAmount = 10 ):
        # code is another function that takes an instruction and context
        # callbacks only take in a CONTEXT object (use Eip to determine instr address)
        def fn( context ):
            hexString = ""
            for i in xrange(0, grabAmount):
                bValue = hex(hilo.readInt(self.hProc, context.Eip + i, 1))[2:]
                if len( bValue ) == 1:
                    bValue = "0" + bValue
                    
                hexString += bValue
            instructions = Decode(context.Eip, hexString.decode('hex'), Decode32Bits)
            curInstr = instructions[0]
            
            code( curInstr, context )
            nextAddress = determineNextInstruction( curInstr, context )
            return nextAddress
        return fn
    def remoteFunction( self, address, argFormat ):
        if ( self.open == False):
            raise Exception()
        def fn(args):
            hThread = hilo.makeRemoteThread( self.hProc )
            cContext = hilo.getThreadContext( hThread )
            # push a 0 onto the stack (for ExitThread)
            cContext.Esp -= 4
            self.writeInt( cContext.Esp, 4, 0 )

            # set next instruction to execute
            cContext.Eip = address
            
            # modify context based on argObject
            for reg in argFormat.regActions:
                setattr( cContext, reg, args[argFormat.regActions[reg]] )

            for (length,name) in argFormat.stackActions:
                cContext.Esp -= length
                self.writeInt( cContext.Esp, length, args[name] )
            # push a return address onto the stack (ExitThread)
            cContext.Esp -= 4
            self.writeInt( cContext.Esp, 4, hilo.getDefaultReturnAddress() )

            hilo.setThreadContext( hThread, cContext )
            hilo.resumeAndWait( hThread )
            
            hilo.closeHandle( hThread )
            
        return fn
    @staticmethod
    def getRunningProcesses( ):
        pass

#
# Defines a way to move argument values in Python to a context for a thread in C++
# Returns calling object to enable call chaining.
#
class ArgFormat:
    def __init__( self ):
    # format for stackActions elements:
    #   tuples: (length,name)
    # * regActions:
    #   associate regName -> argName
        self.stackActions = []
        self.regActions = {}
    def setReg( self, regName, argName ):
        self.regActions[regName] = argName
        return self
    def pushVal( self, length, argName ):
        self.stackActions.append( (length, argName) )
        return self

def determineNextInstruction(instr, context):
    nextAddress = 0
    if ( curInstr[2][0:3] == "RET" ):
        if ( self.readInt( context.Esp, 4 ) != hilo.getDefaultReturnAddress() ):
            nextAddress = self.readInt( context.Esp, 4 )
    else:
        nextAddress = context.Eip + curInstr[1]
    return nextAddress
    
expProg = Process("ExpProg.exe")
DoStuffAddress = 0x00401000
fn = expProg.remoteFunction( DoStuffAddress,
                             ArgFormat().pushVal(4,"doStuff") )
argsToFn = {'doStuff': 7}
def myFun(instr,context):
    print instr
hilo.traceIntoCall( expProg.hProc, DoStuffAddress, expProg.genCallbackTrace(myFun), fn, argsToFn )
expProg.detach()
